package com.fangdd.logtrace.dubbo;

import com.alibaba.dubbo.rpc.*;
import com.fangdd.logtrace.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.Callable;

/**
 * Created by sulong on 16-9-6.
 */
public class LogTraceDubboConsumerFilter implements Filter {
    private static final Logger LOGGER = LoggerFactory.getLogger(LogTraceDubboConsumerFilter.class);

    @Override
    public Result invoke(final Invoker<?> invoker, final Invocation invocation) throws RpcException {
        TraceId next = TraceIdThreadLocal.sequentialNext();
        if (!next.isUnset()) {
            String nextTraceIdStr = next.toString();
            return doInvoke(invoker, invocation, nextTraceIdStr);
        }

        final TraceId newRoot = new TraceId(TraceRootGenerator.newUUIDRoot());
        try {
            return TraceWith.run(newRoot, new Callable<Result>() {
                @Override
                public Result call() throws Exception {
                    TraceId nextId = TraceIdThreadLocal.sequentialNext();
                    return doInvoke(invoker, invocation, nextId.toString());
                }
            });
        } catch (Exception e) {
            LOGGER.debug("unexpected Exception", e);
            throw new RpcException(e);
        }
    }

    private Result doInvoke(Invoker<?> invoker, Invocation invocation, String nextTraceIdStr) {
        RpcContext.getContext().setAttachment(DubboConstants.TRACEID_NAME, nextTraceIdStr);
        long startTime = System.currentTimeMillis();
        String interfaceName = makeInterfaceName(invoker, invocation);
        try {
            BoundsLog.callerIn(LOGGER, nextTraceIdStr, interfaceName);
            Result result = invoker.invoke(invocation);
            BoundsLog.callerOut(LOGGER, nextTraceIdStr, interfaceName, true, System.currentTimeMillis() - startTime);
            return result;
        } catch (RpcException rpcException) {
            if (rpcException.isBiz()) {
                BoundsLog.callerOut(LOGGER, nextTraceIdStr, interfaceName, true, System.currentTimeMillis() - startTime);
            } else {
                BoundsLog.callerOut(LOGGER, nextTraceIdStr, interfaceName, false, System.currentTimeMillis() - startTime);
            }
            throw rpcException;
        } catch (Exception e) {
            BoundsLog.callerOut(LOGGER, nextTraceIdStr, interfaceName, false, System.currentTimeMillis() - startTime);
            throw new RpcException(e);
        }
    }

    public static String makeInterfaceName(Invoker<?> invoker, Invocation invocation) {
        StringBuilder sb = new StringBuilder();
        sb.append("dubbo:");
        if (invoker != null && invoker.getInterface() != null) {
            sb.append(invoker.getInterface().getSimpleName()).append(".");
        }
        sb.append(invocation.getMethodName());
        return sb.toString();
    }
}
